//
// Created by xag on 2020/12/15.
//

#ifndef COTTON_DETECT_COTTON_H
#define COTTON_DETECT_COTTON_H

#include <opencv2/opencv.hpp>
#include "NvInfer.h"


class Cotton2 {
    // struct DetectRes {
    //     int classes;
    //     float x;
    //     float y;
    //     float w;
    //     float h;
    //     float prob;
    // };

public:
    struct DetectRes {
        int classes;
        float x;
        float y;
        float w;
        float h;
        float prob;
    };
    Cotton2(const std::string &config_file);
    ~Cotton2();
    void LoadEngine();
    void InferenceFolder_single2(cv::Mat &src_img, std::map<int, std::vector<int>> &object_number);
    std::map<int, std::vector<int>> EngineInference_single2(cv::Mat &image_name, const int &outSize1, const int &outSize2, void **buffers,
                         const std::vector<int64_t> &bufferSize, cudaStream_t stream);
    
private:
    std::vector<float> prepareImage_single2(cv::Mat &src_img);                                               
    void NmsDetect2(std::vector<DetectRes> &detections);
    float IOUCalculate2(const DetectRes &det_a, const DetectRes &det_b);
    static float sigmoid2(float );
    std::string engine_file;
    std::map<int, std::string>  detect_labels;
    int BATCH_SIZE;
    int INPUT_CHANNEL;
    int IMAGE_WIDTH;
    int IMAGE_HEIGHT;
    nvinfer1::ICudaEngine *engine = nullptr;
    nvinfer1::IExecutionContext *context = nullptr;
    float obj_threshold;
    float nms_threshold;
    int w_heat;
    int h_heat;
};


#endif //COTTON_DETECT_COTTON_H

